import cv2
import numpy as np

class kalman_filter:
    """
    一个带有“门控”机制的鲁棒卡尔曼滤波器。
    它能有效抵抗因严重噪声导致的测量值偏差.
    """
    def __init__(self, p=1000, q=0.01, r=0.5, gating_threshold=100):
        """
        初始化滤波器。

        :param p: 初始后验误差协方差 (P) - 对初始状态的不确定性，可以设大一些。
        :param q: 过程噪声协方差 (Q) - 模型预测的噪声。值越小，轨迹越平滑，但对突变响应慢。
        :param r: 测量噪声协方差 (R) - 测量值的噪声。值越小，越相信测量值。
        :param gating_threshold: 门控阈值（单位：像素）。测量值与预测值的距离超过此值将被视为野值并被丢弃。
        """
        # 1. 初始化OpenCV卡尔曼滤波器
        # 状态向量 [x, y, dx, dy] - (位置x, 位置y, 速度x, 速度y)
        # 测量向量 [x, y]
        self.kf = cv2.KalmanFilter(4, 2)
        
        # 测量矩阵 H: 将状态向量映射到测量向量
        self.kf.measurementMatrix = np.array([[1, 0, 0, 0], [0, 1, 0, 0]], np.float32)
        
        # 转移矩阵 A: 描述状态如何从上一时刻转移到当前时刻
        # x_k = x_{k-1} + dx_{k-1}
        # y_k = y_{k-1} + dy_{k-1}
        self.kf.transitionMatrix = np.array([[1, 0, 1, 0], [0, 1, 0, 1], [0, 0, 1, 0], [0, 0, 0, 1]], np.float32)
        
        # 2. 初始化噪声协方差矩阵
        self.kf.errorCovPost = np.eye(4, dtype=np.float32) * p
        self.kf.processNoiseCov = np.eye(4, dtype=np.float32) * q
        self.kf.measurementNoiseCov = np.eye(2, dtype=np.float32) * r
        
        # 3. 初始化滤波器状态
        self.kf.statePost = np.zeros((4, 1), np.float32)
        
        # 4. 初始化自定义状态变量
        self.gating_threshold = gating_threshold
        self.last_valid_prediction = np.zeros((4, 1), np.float32)
        self.lost_count = 0
        self.max_lost_count = 30  # 连续丢失超过此帧数则认为目标丢失

    def update(self, measurement_point):
        """
        :param measurement_point: 一个包含(x, y)坐标的元组/列表，或者为None（如果目标未被检测到）。
        :return: 滤波后的状态向量 [x, y, dx, dy]^T。
        """
        prediction = self.kf.predict()
        # print("kalman", measurement_point)

        if measurement_point[0] == -1 and measurement_point[1] == -1:
            self.lost_count += 1
            vx = self.kf.statePost[2]
            vy = self.kf.statePost[3]
            speed = np.sqrt(vx**2 + vy**2)

            # 动态调整 decay_factor，速度越大，衰减越小
            min_decay = 0.6
            max_decay = 0.95
            # 线性映射，速度越快 decay_factor 越趋近于 min_decay
            threshold_speed = 50.0
            
            if speed < 1.0:
                decay_factor = max_decay  # 几乎不衰减
            elif speed > 30:
                decay_factor = min_decay
            else:
                norm_speed = min(speed / threshold_speed, 1.0)
                decay_factor = max_decay - (max_decay - min_decay) * (norm_speed) * 1.35

            self.kf.statePost[2] *= decay_factor
            self.kf.statePost[3] *= decay_factor

            if self.lost_count > self.max_lost_count:
                return None
            self.last_valid_prediction = prediction
            return prediction

        measurement = np.array([[np.float32(measurement_point[0])],
                                [np.float32(measurement_point[1])]])

        # 检查是否为初始状态，进行第一次初始化
        if np.all(self.kf.statePost == 0):
            self.kf.statePost[:2, 0] = measurement[:, 0]
            self.kf.statePost[2:, 0] = 0  # 初始速度为0
            self.last_valid_prediction = self.kf.statePost.copy()
            self.lost_count = 0
            return self.kf.statePost

        predicted_pos = prediction[:2].astype(np.float32)
        distance = cv2.norm(measurement, predicted_pos)
        # print(f"[{measurement_point}] <-> predict: {predicted_pos.ravel()} | distance: {distance}")

        if distance < self.gating_threshold:
            self.lost_count = 0
            corrected_state = self.kf.correct(measurement)
            self.last_valid_prediction = corrected_state
            return corrected_state
        else:
            self.lost_count += 1
            print(f"测量值 {measurement_point} 被拒绝 (野值!) -> 输出预测值: {prediction.ravel()}")
            if self.lost_count > self.max_lost_count:
                self.reset_filter()
                return None
            return prediction


    def get_state(self):
        """返回当前最新的有效状态"""
        return self.last_valid_prediction

    def is_lost(self):
        """判断目标是否已确认丢失"""
        return self.lost_count > self.max_lost_count
    
    def reset_filter(self):
        """重置卡尔曼滤波器到初始状态"""
        self.kf.statePost = np.zeros((4, 1), np.float32)
        self.last_valid_prediction = np.zeros((4, 1), np.float32)
        self.lost_count = 0

    def get_velocity(self):
        """返回当前速度 (dx, dy)"""
        # dx = self.last_valid_prediction[2].astype(np.float32)
        # dy = self.last_valid_prediction[3].astype(np.float32)
        dx = float(self.last_valid_prediction[2])
        dy = float(self.last_valid_prediction[3])
        return dx, dy
    


# # ================== 示例用法 ==================

# if __name__ == '__main__':
#     # 1. 初始化滤波器
#     # 假设一个测量点，如果与预测点距离超过30个像素，就被认为是噪声

#     kf = kalmanFilter(gating_threshold=30)

#     # # 2. 模拟一系列测量值
#     # true_points = [(10, 10), (15, 15), (20, 20), (25, 25), (30, 30)]
#     # 在第6个点引入一个严重的噪声偏差
#     measurements = [
#         (11, 9), (14, 16), (21, 19), (26, 24), (29, 31),
#         (150, 150),  # <--- 这是一个严重的“野值”
#         (35, 35), (40, 40),
#         None,        # <--- 模拟一次目标丢失
#         (50, 50)
#     ]

#     print("开始模拟...")
#     print("格式: [测量值] -> [滤波器输出位置(x, y)]")
#     print("-" * 50)

#     for i, meas in enumerate(measurements):
#         # 3. 更新滤波器
#         state = kf.update(meas)

#         if state is not None:
#             output_pos = (state[0, 0], state[1, 0])
#             # === FIX: Convert tuples to numpy arrays for cv2.norm ===

#             if meas is not None and cv2.norm(np.array(meas, dtype=np.float32), np.array(output_pos)) > kf.gating_threshold:
#                  print(f"第{i+1}帧: 测量值 {meas} 被拒绝 (野值!) -> 输出预测值: ({output_pos[0]:.2f}, {output_pos[1]:.2f})")

#             else:
#                  print(f"第{i+1}帧: 测量值 {meas} -> 输出: ({output_pos[0]:.2f}, {output_pos[1]:.2f})")

#         else:
#             if kf.is_lost():
#                 print(f"第{i+1}帧: 测量值 {meas} -> 目标已确认丢失!")
#             else:
#                  print(f"第{i+1}帧: 测量值 {meas} -> 目标丢失，使用预测值")



#     # 预期结果：
#     # 前5帧，输出会平滑地跟随测量值。
#     # 第6帧，测量值(150, 150)会因距离过大而被拒绝，滤波器会输出一个基于(30,30)的预测值（大约是35,35）。
#     # 第7、8帧，滤波器会从(35,35)的状态继续平滑地跟随新的正确测量值。
#     # 第9帧，目标丢失，滤波器会输出一个预测值（大约是45,45）。
#     # 第10帧，滤波器会从(45,45)的状态继续跟随。
